package com.liuqi.openai.model.common;

import com.liuqi.openai.core.chat.ChatCompletionChoice;
import com.liuqi.openai.core.chat.ChatCompletionResponse;
import com.liuqi.openai.core.chat.Delta;
import com.liuqi.openai.core.chat.Message;
import com.liuqi.openai.core.chat.Role;
import com.liuqi.openai.core.chat.ToolCall;
import com.liuqi.openai.core.client.OpenAiClient;
import com.liuqi.openai.core.shared.Usage;
import com.liuqi.openai.data.message.AiMessage;
import com.liuqi.openai.data.message.ChatMessage;
import com.liuqi.openai.data.message.ToolResultMessage;
import com.liuqi.openai.data.message.UserMessage;
import com.liuqi.openai.model.output.FinishReason;
import com.liuqi.openai.model.output.TokenUsage;
import com.liuqi.openai.tool.ToolCallRequest;
import java.util.List;

import static com.liuqi.openai.model.output.FinishReason.CONTENT_FILTER;
import static com.liuqi.openai.model.output.FinishReason.LENGTH;
import static com.liuqi.openai.model.output.FinishReason.STOP;
import static com.liuqi.openai.model.output.FinishReason.TOOL_CALLS;
import static com.liuqi.openai.util.ValidationUtil.ensureNotBlank;
import static com.liuqi.openai.util.ValidationUtil.ensureNotNull;
import static java.util.stream.Collectors.toList;

/**
 *
 * openai 模型, 公共抽象的类
 *
 * @author liuqi
 * @date 2025/7/19
 **/
public abstract class AbstractOpenAiModel {

    /**
     * openai 客户端
     *
     */
    protected final OpenAiClient openAiClient;

    /**
     * 模型名称
     */
    protected final String modelName;

    protected AbstractOpenAiModel(AbstractOpenAiModelBuilder builder) {
        this.modelName = ensureNotBlank(builder.modelName, "modelName");
        this.openAiClient = ensureNotNull(builder.openAiClient, "openAiClient");
    }

    /**
     * 构建 TokenUsage
     *
     * @param usage
     * @return
     */
    protected TokenUsage tokenUsageFrom(Usage usage) {
        if (usage == null) {
            return null;
        }
        return new TokenUsage(
                usage.getPromptTokens(),
                usage.getCompletionTokens(),
                usage.getTotalTokens()
        );
    }

    /**
     * 构建 FinishReason
     *
     * @param finishReason
     * @return
     */
    protected FinishReason finishReasonFrom(String finishReason) {
        switch (finishReason) {
            case "stop":
                return STOP;
            case "length":
                return LENGTH;
            case "tool_calls":
                return TOOL_CALLS;
            case "content_filter":
                return CONTENT_FILTER;
            default:
                return null;
        }
    }

    /**
     * 转换成openai的消息
     *
     * @param messages
     * @return
     */
    protected List<Message> toOpenAiMessages(List<ChatMessage> messages) {
        return messages.stream()
                .map(this::toOpenAiMessage)
                .collect(toList());
    }

    /**
     * 转换成openai的消息
     *
     * @param message
     * @return
     */
    protected Message toOpenAiMessage(ChatMessage message) {
        return Message.builder()
                .role(roleFrom(message))
                .name(nameFrom(message))
                .content(message.text())
                .toolCallId(message instanceof ToolResultMessage ? ((ToolResultMessage) message).toolCallId() : null)
                .toolCalls(toolCallsFrom(message))
                .build();
    }

    /**
     * 转换成 AiMessage 消息
     *
     * @param response
     * @return
     */
    protected AiMessage aiMessageFrom(ChatCompletionResponse response) {
        String id = response.getId();
        String reasoningContent = null;
        String content = null;
        List<ToolCallRequest> toolCallRequests = null;

        ChatCompletionChoice choice = response.getChoices().get(0);
        // 流式响应
        if (choice.getDelta() != null) {
            Delta delta = choice.getDelta();
            reasoningContent = delta.getReasoningContent();
            content = delta.getContent();
            List<ToolCall> toolCalls = delta.getToolCalls();
            if (toolCalls != null) {
                toolCallRequests = toolCallsToToolCallRequests(toolCalls);
            }
        }
        // 非流式响应
        else if (choice.getMessage() != null) {
            content = choice.getMessage().getContent();
            List<ToolCall> toolCalls = choice.getMessage().getToolCalls();
            if (toolCalls != null) {
                toolCallRequests = toolCallsToToolCallRequests(toolCalls);
            }
        }

        return AiMessage.aiMessage(id, content, reasoningContent, toolCallRequests);
    }

    /**
     * ToolCall -> ToolCallRequest
     *
     * @param toolCalls
     * @return
     */
    List<ToolCallRequest> toolCallsToToolCallRequests(List<ToolCall> toolCalls) {
        return toolCalls.stream()
                .map(toolCall -> ToolCallRequest.builder()
                        .index(toolCall.getIndex())
                        .id(toolCall.getId())
                        .type(toolCall.getType())
                        .name(toolCall.getFunction().getName())
                        .arguments(toolCall.getFunction().getArguments())
                        .build()
                )
                .collect(toList());
    }

    /**
     * 消息角色
     *
     * @param message
     * @return
     */
    Role roleFrom(ChatMessage message) {
        switch (message.type()) {
            case SYSTEM:
                return Role.SYSTEM;
            case ASSISTANT:
                return Role.ASSISTANT;
            case USER:
                return Role.USER;
            case TOOL_RESULT:
                return Role.TOOL;
        }
        return Role.USER;
    }

    /**
     * 消息 name
     *
     * @param message
     * @return
     */
    String nameFrom(ChatMessage message) {
        if (message instanceof UserMessage) {
            return ((UserMessage) message).name();
        }
        if (message instanceof ToolResultMessage) {
            return ((ToolResultMessage) message).name();
        }
        return null;
    }

    /**
     * 工具调用
     *
     * @param message
     * @return
     */
    List<ToolCall> toolCallsFrom(ChatMessage message) {
        if (message instanceof AiMessage) {
            AiMessage aiMessage = (AiMessage) message;
            if (aiMessage.toolCallRequests() != null) {
                List<ToolCallRequest> calls = aiMessage.toolCallRequests();
                return calls.stream()
                        .map(request -> ToolCall.builder()
                                // 当tool_calls只有一条, 不设置index, 符合 pydantic
                                .index(calls.size() > 1 ? request.getIndex() : null)
                                .id(request.getId())
                                .type(request.getType())
                                .function(ToolCall.Function.builder().name(request.getName()).arguments(request.getArguments()).build())
                                .build()
                        )
                        .collect(toList());
            }
        }
        return null;
    }

    /**
     * openai 模型 构造器
     *
     * @param <B>   泛型: Ai Model {@link AbstractOpenAiModel}
     * @param <T>   泛型: 构造器 {@link AbstractOpenAiModelBuilder}
     */
    public static abstract class AbstractOpenAiModelBuilder<B extends AbstractOpenAiModel, T extends AbstractOpenAiModelBuilder<B, T>> {
        private OpenAiClient openAiClient;
        private String modelName;

        public T openAiClient(OpenAiClient openAiClient) {
            this.openAiClient = openAiClient;
            return (T) this;
        }
        public T modelName(String modelName) {
            this.modelName = modelName;
            return (T) this;
        }
        /* 由子类实现 */
        public abstract B build();
    }

}
